# predict_model.py
import os
import torch
import torch.nn as nn
from torchvision import transforms, models
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import pandas as pd
from tqdm import tqdm
import argparse

# 允许加载截断图像（与训练一致）
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

class TestDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.images = sorted([
            f for f in os.listdir(root_dir)
            if f.lower().endswith(('.png', '.jpg', '.jpeg'))
        ])

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_name = self.images[idx]
        img_path = os.path.join(self.root_dir, img_name)
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, img_name

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--test_dir', type=str, required=True, help='Directory containing test images')
    parser.add_argument('--checkpoint', type=str, required=True, help='Path to model .pth checkpoint')
    parser.add_argument('--num_classes', type=int, default=400)
    parser.add_argument('--output_csv', type=str, default='predictions.csv')
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--num_workers', type=int, default=4)
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
    return parser.parse_args()

def get_transform():
    return transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

def build_model(model_name, num_classes):
    if model_name == 'resnet50':
        model = models.resnet50(weights=None)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
    elif model_name == 'convnext_v2_large':
        model = models.convnext_large(weights=None)
        model.classifier[2] = nn.Linear(model.classifier[2].in_features, num_classes)
    else:
        raise ValueError(f"Unsupported model: {model_name}")
    return model
    
def main():
    args = parse_args()
    device = torch.device(args.device)
    print(f"Using device: {device}")

    # 加载 checkpoint
    checkpoint = torch.load(args.checkpoint, map_location=device)
    model_name = checkpoint.get('model_name', 'resnet50')
    num_classes = checkpoint.get('num_classes', 400)    

    print(f"Loading model: {model_name} with {num_classes} classes")

    # 构建模型（结构必须与训练时一致）
    model = build_model(model_name, num_classes)
    if 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
    else:
        state_dict = checkpoint
    model.load_state_dict(state_dict)    
    model = model.to(device)
    model.eval()

    # 测试数据集
    transform = get_transform()
    test_dataset = TestDataset(root_dir=args.test_dir, transform=transform)
    test_loader = DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True
    )

    # 预测
    results = []
    with torch.no_grad():
        for inputs, filenames in tqdm(test_loader, desc="Predicting"):
            inputs = inputs.to(device)
            outputs = model(inputs)
            _, preds = outputs.max(1)
            for fname, pred in zip(filenames, preds.cpu().numpy()):
                class_str = f"{pred:04d}"  # 补零到四位
                results.append((fname, class_str))

    # 自动创建父目录
    output_path = args.output_csv
    os.makedirs(os.path.dirname(output_path), exist_ok=True)  

    # 保存 CSV（无 header，逗号分隔）
    df = pd.DataFrame(results, columns=['filename', 'class'])    
    df.to_csv(output_path, index=False, header=False)
    print(f"Predictions saved to {output_path}")

if __name__ == '__main__':
    main()